In the decode path with rmsnorm=True, the previous code passed z=None to selective_state_update and then applied the gate and normalization together in the gated RMSNorm kernel:

y = selective_state_update(..., z=None, ...)       # SSM writes y
y = self.norm(y, z)                                # silu(z)*y, rmsnorm, weight

selective_state_update already supports applying silu(z)*y inline via its HAS_Z path -- the work then happens in fp32 registers right after the state-C reduction, before y's bf16 round-trip out to HBM. Swapping the gate site lets the downstream RMSNormGated take its z=None fast path (skipping the gate step entirely):

y = selective_state_update(..., z=z_reshaped, ...) # SSM writes silu(z)*y
y = self.norm(y, None)                             # rmsnorm, weight

Net effect:

• One fewer HBM round-trip of z (SSM reads z in-kernel instead of the post-SSM gated-norm reading y and z separately).
• Cheaper per-call cost for the gated norm (no gate work).
• Math is identical up to bf16 rounding; in fact slightly more precise because y no longer round-trips through bf16 between SSM and norm.

Measured on nano-v3 at BS=1, OSL=256, 10 iterations, outlier-trimmed p50:

• gate_shift_off: 255.0 tok/s p50
• gate_shift_on: 257.5 tok/s p50 (+1.0%)

What does this PR do ?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.